#!/usr/bin/python

import logging
import random

from hips.scheduling import constraint
from hips.scheduling import solver
from hips.scheduling import solution
from hips.scheduling import solution_log


class Domain(list):
  """Class used to control possible values for variables

  When list or tuples are used as domains, they are automatically
  converted to an instance of that class.
  """

  def __init__(self, set):
    """
    @param set: Set of values that the given variables may assume
    @type  set: set of objects comparable by equality
    """
    list.__init__(self, set)
    self._hidden = []
    self._states = []

  def ResetState(self):
    """Reset to the original domain state, including all possible values."""
    self.extend(self._hidden)
    del self._hidden[:]
    del self._states[:]

  def PushState(self):
    """Save current domain state.
    
    Variables hidden after that call are restored when that state
    is popped from the stack.
    """
    self._states.append(len(self))

  def PopState(self):
    """Restore domain state from the top of the stack.

    Variables hidden since the last popped state are then available
    again.
    """
    diff = self._states.pop()-len(self)
    if diff:
      self.extend(self._hidden[-diff:])
      del self._hidden[-diff:]

  def HideValue(self, value):
    """Hide the given value from the domain.

    After that call the given value won't be seen as a possible value
    on that domain anymore. The hidden value will be restored when the
    previous saved state is popped.

    @param value: Object currently available in the domain
    """
    list.remove(self, value)
    self._hidden.append(value)


class SchedulingProblem(object):

  def __init__(self, schedule):
    self.doctors = []
    self.schedule = schedule
    # TODO Encode counts in the schedule spec?
    self.day_count = 2
    self.night_count = 1
    self.solutions = None
    self._variables = {}

  def AddDoctor(self, doc):
    self.doctors.append(doc)

  def CreateVariables(self, doc_count, day_shifts, night_shifts):
    domains = {}

    # Day shift values
    day_shift_values = []
    for k in range(doc_count):
      for m in range(k):
        day_shift_values.append((1 << k) | (1 << m))
    for v in day_shifts:
      domains[v] = Domain(day_shift_values)

    # Night shift values
    night_shift_values = [(1 << i) for i in range(doc_count)]
    for v in night_shifts:
      domains[v] = Domain(night_shift_values)

    return domains

  def CreateConstraints(self):
    r = self.schedule.AllShiftKeysArray()
    shifts = self.schedule.ShiftCount()
    doc_count = len(self.doctors)
    constraints = []

    # Global constraints
    self.log = solution_log.CoarseSolutionLog()
    constraints.append(
        constraint.NoTwoConsecutiveShifts(doc_count, self.schedule, self.log))
    constraints.append(
        constraint.NoLessThanMinConsecutiveDayShifts(
            doc_count, self.schedule, self.log))
    constraints.append(
        constraint.DayShiftsOverlap(doc_count, self.schedule, self.log))

    # Per-doc constraints
    min_shifts = []
    max_shifts = []
    max_night_shifts = []
    consecutive_days = []
    for i, doc in enumerate(self.doctors):
      doc_mask = 1 << i
      constraints.append(constraint.NotInDay(
          doc.days_off, doc_mask, doc_count, self.schedule, self.log))
      min_shifts.append(doc.min_shifts)
      max_shifts.append(doc.max_shifts)
      max_night_shifts.append(doc.max_night_shifts)
      consecutive_days.append(doc.max_consecutive_days)
    constraints.append(constraint.TotalShiftsInRange(
        min_shifts, max_shifts, max_night_shifts, doc_count,
        self.schedule, self.log))
    constraints.append(constraint.NoMoreThanMaxConsecutiveDays(
        consecutive_days, doc_count, self.schedule, self.log))
    return constraints

  def Solve(self, steps=1000):
    s = solver.InitialConditionsMinConflictsSolver(
        self.doctors, self.schedule, steps=steps)

    domains = self.CreateVariables(len(self.doctors),
        self.schedule.DayShiftKeysArray(),
        self.schedule.NightShiftKeysArray())
    all_variables = domains.keys()

    constraints = []
    for constraint in self.CreateConstraints():
      constraints.append((constraint, all_variables))

    vconstraints = {}
    for variable in domains:
      vconstraints[variable] = []

    for constraint, variables in constraints:
      for variable in variables:
        vconstraints[variable].append((constraint, variables))

    for constraint, variables in constraints[:]:
      constraint.PreProcess(variables, domains, constraints, vconstraints)

    for domain in domains.values():
      domain.ResetState()
      if not domain:
        return None, None, None

    raw_solution = s.Solve(domains, constraints, vconstraints)
    self.solution = solution.Solution(self.doctors, self.schedule, raw_solution)
    self.iterations = s.Iterations()

    return self.solution

  def GetSolution(self):
    return self.solution

  def GetSolutionLog(self):
    return self.log
